package top.yudoge.dogeeditor.llm.rewrite

import cn.hutool.core.util.StrUtil
import dev.langchain4j.memory.ChatMemory
import dev.langchain4j.memory.chat.MessageWindowChatMemory
import dev.langchain4j.model.chat.ChatLanguageModel
import dev.langchain4j.model.chat.StreamingChatLanguageModel
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.stereotype.Component
import top.yudoge.dogeeditor.llm.config.models.ModelRegistry
import top.yudoge.dogeeditor.llm.exceptions.RewriteException
import top.yudoge.dogeeditor.llm.memory.SessionMemory
import top.yudoge.dogeeditor.llm.objects.query.RewriteQuery


/**
 * @author Yudoge
 * @since 2025/04/13
 *
 * SQL Rewriter take a user SQL code snippet and a task, using pre-training data(RAG) and real-time database 
 * agent to help user to rewrite their code. If the user code is empty, the SQL Rewriter will generate
 * a new SQL code in scratch.
 *
 * SQL Rewriter will start multiple round conversation to llm to complete the task. 
 *
 * Design:
 *  1. extract table name from user code + user input, and guessing possible table name from user task.
 *  2. get table information using these table name
 *  3. using the user task to match related information(RAG)
 *  4. generate sql using the information above
 *  5. self checking if the generated sql is the `select` sql
 *     5.1. pass: return the sql
 *     5.2. fail: give the error msg to llm and regenerate
 *  
 *  Each step will be designed as a chain node of a sequential chain and they share the same context(like memory, user-callback etc.)
 */

@Component
class SQLRewriter {

    @Autowired
    private lateinit var sessionMemory: SessionMemory

    @Autowired
    private lateinit var modelRegistry: ModelRegistry


    @Autowired
    private lateinit var tableNameExtractChainNode: TableNameExtractChainNode

    @Autowired
    private lateinit var tableInformationGetChainNode: TableInformationGetChainNode

    @Autowired
    private lateinit var ragChainNode: RAGChainNode

    @Autowired
    private lateinit var SQLGenerateChainNode: SQLGenerateChainNode

    @Autowired
    private lateinit var SQLSelfCheckingChainNode: SQLSelfCheckingChainNode

    @Autowired
    private lateinit var getPossibleTableChainNode: GetPossibleTableChainNode

    @Autowired
    private lateinit var sqlGenerate2ChainNode: SQLGenerate2ChainNode

    @Autowired
    private lateinit var userInputAnalyseChainNode: UserInputAnalyseChainNode

    @Autowired
    private lateinit var userQueryRewriteChainNode: UserQueryRewriteChainNode


    fun rewrite(query: RewriteQuery, callback: RewriteCallback) {
        if (!preCheckQuery(query, callback)) {
            return
        }

        val streamingModel = modelRegistry.getStreamingChatModel(query.modelName)
        val chatModel = modelRegistry.getChatModel(query.modelName)

        TaskChain(
            nodes = listOf(
                userInputAnalyseChainNode,
                getPossibleTableChainNode,
                sqlGenerate2ChainNode

//                tableNameExtractChainNode,            // step1: extract table name from user code + user input, and guessing possible table name from user task.
//                tableInformationGetChainNode,         // step2: get table information using these table name
//                ragChainNode,                         // step3: using the user task to match related information(RAG)
//                SQLGenerateChainNode,                 // step4: generate sql using the information above
//                SQLSelfCheckingChainNode,             // step5: self checking if the generated sql is the `select` sql
            )
        ).apply {
            process(query, callback, streamingModel, chatModel)
        }
    }

    fun preCheckQuery(query: RewriteQuery, callback: RewriteCallback): Boolean {
        // 0. check the query is valid
        if (!StrUtil.isAllNotBlank(query.sessionId, query.userTask, query.dataSourceIdentifier, query.databaseName, query.schemaName)) {
            callback.onEvent(RewriteEvent(
                RewriteEventType.Error,
                RewriteException("query is not valid")
            ))
            return false
        }

        if (StrUtil.isBlank(query.modelName) || !modelRegistry.hasChatModel(query.modelName) || !modelRegistry.hasStreamingChatModel(query.modelName)) {
            callback.onEvent(RewriteEvent(
                RewriteEventType.Error,
                RewriteException("model ${query.modelName} is not found")
            ))
            return false
        }

        // 1. check the session is created
        if (!sessionMemory.containsSession(query.sessionId)) {
            callback.onEvent(RewriteEvent(
                RewriteEventType.Error,
                RewriteException("session ${query.sessionId} is not found")
            ))
            return false
        }

        return true;
    }
    
}